#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
ComfyUI LoRA训练插件依赖安装脚本
解决huggingface_hub版本兼容性问题
"""

import os
import sys
import subprocess
import importlib

def run_command(command, description=""):
    """运行命令并显示结果"""
    print(f"🔄 {description}")
    print(f"执行命令: {command}")
    
    try:
        result = subprocess.run(command, shell=True, capture_output=True, text=True)
        if result.returncode == 0:
            print(f"✅ {description} 成功")
            if result.stdout:
                print(f"输出: {result.stdout.strip()}")
        else:
            print(f"❌ {description} 失败")
            if result.stderr:
                print(f"错误: {result.stderr.strip()}")
        return result.returncode == 0
    except Exception as e:
        print(f"❌ {description} 异常: {e}")
        return False

def check_python_version():
    """检查Python版本"""
    version = sys.version_info
    print(f"🐍 Python版本: {version.major}.{version.minor}.{version.micro}")
    
    if version.major < 3 or (version.major == 3 and version.minor < 8):
        print("❌ Python版本过低，需要Python 3.8或更高版本")
        return False
    
    print("✅ Python版本符合要求")
    return True

def check_huggingface_hub():
    """检查huggingface_hub版本"""
    try:
        import huggingface_hub
        version = huggingface_hub.__version__
        print(f"📦 huggingface_hub版本: {version}")
        
        # 检查cached_download是否可用
        try:
            from huggingface_hub import cached_download
            print("✅ cached_download可用")
            return True
        except ImportError:
            print("❌ cached_download不可用，需要修复")
            return False
            
    except ImportError:
        print("❌ huggingface_hub未安装")
        return False

def fix_huggingface_hub():
    """修复huggingface_hub版本问题"""
    print("🔧 开始修复huggingface_hub版本问题...")
    
    # 卸载当前版本
    if not run_command(f"{sys.executable} -m pip uninstall huggingface_hub -y", "卸载当前huggingface_hub版本"):
        print("⚠️ 卸载失败，继续尝试安装兼容版本")
    
    # 安装兼容版本
    if run_command(f"{sys.executable} -m pip install huggingface_hub==0.24.5", "安装huggingface_hub==0.24.5"):
        # 验证安装
        if check_huggingface_hub():
            print("✅ huggingface_hub修复成功")
            return True
    
    print("❌ huggingface_hub修复失败")
    return False

def install_requirements():
    """安装requirements.txt中的依赖"""
    print("📦 安装项目依赖...")
    
    requirements_file = os.path.join(os.path.dirname(__file__), "requirements.txt")
    if not os.path.exists(requirements_file):
        print(f"❌ requirements.txt文件不存在: {requirements_file}")
        return False
    
    if run_command(f"{sys.executable} -m pip install -r {requirements_file}", "安装项目依赖"):
        print("✅ 项目依赖安装成功")
        return True
    
    print("❌ 项目依赖安装失败")
    return False

def install_sd_scripts_requirements():
    """安装sd_scripts的依赖"""
    print("📦 安装sd_scripts依赖...")
    
    sd_scripts_req = os.path.join(os.path.dirname(__file__), "sd_scripts", "requirements.txt")
    if not os.path.exists(sd_scripts_req):
        print(f"❌ sd_scripts requirements.txt文件不存在: {sd_scripts_req}")
        return False
    
    if run_command(f"{sys.executable} -m pip install -r {sd_scripts_req}", "安装sd_scripts依赖"):
        print("✅ sd_scripts依赖安装成功")
        return True
    
    print("❌ sd_scripts依赖安装失败")
    return False

def main():
    """主函数"""
    print("🚀 ComfyUI LoRA训练插件依赖安装脚本")
    print("=" * 50)
    
    # 检查Python版本
    if not check_python_version():
        return False
    
    # 检查当前huggingface_hub状态
    if not check_huggingface_hub():
        # 修复huggingface_hub
        if not fix_huggingface_hub():
            print("❌ 无法修复huggingface_hub，请手动安装兼容版本")
            return False
    
    # 安装项目依赖
    if not install_requirements():
        print("⚠️ 项目依赖安装失败，尝试安装sd_scripts依赖")
    
    # 安装sd_scripts依赖
    if not install_sd_scripts_requirements():
        print("⚠️ sd_scripts依赖安装失败")
    
    print("=" * 50)
    print("✅ 依赖安装完成！")
    print("💡 如果仍有问题，请尝试以下命令：")
    print(f"   {sys.executable} -m pip install huggingface_hub==0.24.5")
    print(f"   {sys.executable} -m pip install -r requirements.txt")
    
    return True

if __name__ == "__main__":
    main() 